"""
基于 LangChain 的智能体协作流程调度器
"""
from langchain_core.runnables import RunnablePassthrough, RunnableLambda
from langchain_core.prompts import ChatPromptTemplate
from agents.data_collector_agent import DataCollectorAgent
from agents.data_analyst_agent import DataAnalystAgent
from agents.decision_maker_agent import DecisionMakerAgent
from colorama import Fore, Style
from typing import Dict, Any
import time


class LangChainAgentFlow:
    """基于 LangChain 的智能体协作流程"""
    
    def __init__(self):
        print(Fore.CYAN + "🚀 初始化 LangChain 智能体协作系统...")
        self.data_collector = DataCollectorAgent()
        self.data_analyst = DataAnalystAgent()
        self.decision_maker = DecisionMakerAgent()
        self.workflow_chain = self._create_workflow_chain()
        print(Fore.GREEN + "✅ LangChain 智能体系统初始化完成")

    def _create_workflow_chain(self):
        """创建 LangChain 工作流链"""
        
        def step1_collect_data(inputs: Dict[str, Any]) -> Dict[str, Any]:
            """步骤1：数据收集"""
            stock_code = inputs["stock_code"]
            print(Fore.YELLOW + f"\n🔍 [LangChain 工作流] 步骤 1/3: 数据收集 ({stock_code})")
            
            raw_data = self.data_collector.execute(stock_code)
            
            return {
                "stock_code": stock_code,
                "raw_data_report": raw_data,
                "step1_status": "completed" if not raw_data.startswith("Error") else "failed"
            }
        
        def step2_analyze_data(data: Dict[str, Any]) -> Dict[str, Any]:
            """步骤2：数据分析"""
            if data["step1_status"] == "failed":
                return {**data, "step2_status": "skipped", "analysis_report": "数据收集失败，跳过分析"}
            
            print(Fore.YELLOW + f"\n📊 [LangChain 工作流] 步骤 2/3: 数据分析")
            
            analysis_report = self.data_analyst.execute(data["raw_data_report"])
            
            return {
                **data,
                "analysis_report": analysis_report,
                "step2_status": "completed" if not analysis_report.startswith("Error") else "failed"
            }
        
        def step3_make_decision(data: Dict[str, Any]) -> Dict[str, Any]:
            """步骤3：投资决策"""
            if data["step2_status"] in ["failed", "skipped"]:
                return {**data, "step3_status": "skipped", "decision_report": "分析失败，跳过决策"}
            
            print(Fore.YELLOW + f"\n🎯 [LangChain 工作流] 步骤 3/3: 投资决策")
            
            decision_report = self.decision_maker.execute(data["analysis_report"])
            
            return {
                **data,
                "decision_report": decision_report,
                "step3_status": "completed" if not decision_report.startswith("Error") else "failed"
            }
        
        # 创建 LangChain 工作流链
        workflow_chain = (
            RunnableLambda(step1_collect_data)
            | RunnableLambda(step2_analyze_data)
            | RunnableLambda(step3_make_decision)
        )
        
        return workflow_chain

    def run(self, stock_code: str):
        """执行完整的 LangChain 智能体协作流程"""
        print(Fore.CYAN + Style.BRIGHT + f"\n🎬 === LangChain 智能体协作开始 ===")
        print(Fore.CYAN + f"📈 目标股票: {stock_code}")
        print(Fore.CYAN + f"🤖 智能体数量: 3 个")
        print(Fore.CYAN + f"⚡ 处理引擎: LangChain")
        
        start_time = time.time()
        
        try:
            # 使用 LangChain 工作流链执行
            print(Fore.MAGENTA + f"\n🔄 [LangChain] 启动智能体协作工作流...")
            
            result = self.workflow_chain.invoke({"stock_code": stock_code})
            
            end_time = time.time()
            processing_time = round(end_time - start_time, 2)
            
            # 显示结果
            self._display_results(result, processing_time)
            
        except Exception as e:
            print(Fore.RED + f"\n❌ LangChain 工作流执行失败: {str(e)}")

    def _display_results(self, result: Dict[str, Any], processing_time: float):
        """显示处理结果"""
        print(Fore.GREEN + f"\n✅ [LangChain] 智能体协作工作流完成 (耗时: {processing_time}s)")
        
        # 显示各步骤状态
        print(Fore.CYAN + "\n📋 === 执行状态摘要 ===")
        steps = [
            ("数据收集", result.get("step1_status", "unknown")),
            ("数据分析", result.get("step2_status", "unknown")),
            ("投资决策", result.get("step3_status", "unknown"))
        ]
        
        for step_name, status in steps:
            status_color = Fore.GREEN if status == "completed" else Fore.YELLOW if status == "skipped" else Fore.RED
            status_icon = "✅" if status == "completed" else "⏭️" if status == "skipped" else "❌"
            print(f"{status_color}{status_icon} {step_name}: {status}")
        
        # 显示分析报告
        if result.get("step2_status") == "completed":
            print(Fore.WHITE + Style.BRIGHT + "\n📊 === 数据分析报告 ===")
            print(Fore.WHITE + result["analysis_report"][:500] + "..." if len(result["analysis_report"]) > 500 else result["analysis_report"])
        
        # 显示最终决策
        if result.get("step3_status") == "completed":
            print(Fore.CYAN + Style.BRIGHT + "\n🎯 === 最终投资建议 ===")
            print(Fore.CYAN + result["decision_report"])
        
        print(Fore.CYAN + Style.BRIGHT + f"\n🏁 === LangChain 智能体协作结束 ===")

    def get_agent_info(self):
        """获取智能体信息"""
        agents_info = [
            self.data_collector.get_agent_info(),
            self.data_analyst.get_agent_info(),
            self.decision_maker.get_agent_info()
        ]
        
        print(Fore.CYAN + "\n🤖 === LangChain 智能体信息 ===")
        for i, info in enumerate(agents_info, 1):
            print(f"{Fore.YELLOW}{i}. {info['name']}")
            print(f"   - 类型: {info['agent_type']}")
            print(f"   - 模型: {info['model']}")
            print(f"   - 服务: {info['base_url']}")
        
        return agents_info


# 保持向后兼容性
class AgentFlow(LangChainAgentFlow):
    """向后兼容的智能体流程类"""
    pass

